#include "densityofstates.h"


DensityOfStates::DensityOfStates(void) : 
	// A pointer to the used system. The system contains variables 
	// needed to construct elements (coupling strengths etc.).
	system(&System::Std),
	// A double containing the current bias (for bi-layered 
	// systems).
	bias(0.),
	// A  double vector containing the energy sample points used in
	// the calculations. The size of this has little influence on 
	// the operation time.
	E(arma::zeros<vec>(1)),
	// A double containing the broadening factor.
	eta(0.001),
	// An integer containing the number of k-points to use in the
	// calculations of the Green's function matrix.
	N_samples(1000)
{ }

// This operator (use by simply appending () to a DensityOfStates
// type variable) calculates the density of states for the unit
// id. The operator returns the object it self for easy access
// to result. The result is stored in the DoS result variable. The 
// variable may be retrieved immidiately after calculation i.e:
// dos(id).DoS // Calculate and retreive result.
DensityOfStates & DensityOfStates::operator()(int id)
{
	// First we get a set of sample k-points.
	int weight; // the weight of each k-point!
	mat k = getPseudoUniformKDistribution(id, weight);

	// An identity matrix!
	const mat I = arma::eye<mat>(system->getUnit(id)->getAtoms().size(),system->getUnit(id)->getAtoms().size());
	const complex ii = complex(0,1);

	// The update step is calculated! Note that updateProgress(step)
	// must be called N times = 1/step. This will give a full update 
	// of 1!
	step = 1./(double)(k.n_cols*E.n_elem);
	
	// The dos value is allocated!
	mat fullDoS(E.n_elem,k.n_cols);

	_Cilk_for(int nk = 0; nk < k.n_cols; nk++)
	{
		// Get the Hamilton!
		cx_mat H = getHamiltonMatrix(id, k.col(nk));

		for(int nE = 0; nE < E.n_elem; nE++)
		{
			cx_mat tmpG = arma::inv(I*complex(E(nE),eta) - H);

			fullDoS(nE,nk) = arma::trace( ii*(tmpG-tmpG.ht())).real() / 2. / arma::datum::pi / k.n_cols;

			// The progress is updated
			try_updateProgress();
		}
	}
	
	DoS = arma::sum(fullDoS,1);

	return *this;
}

// This function attempts to update the progress.
void DensityOfStates::try_updateProgress() const
{
	if(!updateProgress)
		return;

	updateProgress(step);
}


// This function returns the Hamilton matrix H.
cx_mat DensityOfStates::getHamiltonMatrix(int id, const vec &k) const
{
	// Make a Hamiton object and set the options.
	Hamiltonian hamilton;
	hamilton.system = system;
	hamilton.bias = bias;
	hamilton.k = k;

	// Calculate and return the dense matrix.
	return cx_mat(hamilton(id).H);
}

// A random point generator (in triangle)!
vec DensityOfStates::randomPoint(const vec &p_1, const vec &p_2, const vec &p_3) {
	vec2 f = arma::randu<vec>(2);
	return (1-sqrt(f(1)))*p_3 + sqrt(f(1))*((1-f(0))*p_1+f(0)*p_2);
}

// This function returns a randomly generated uniform distribution of k-points.
mat DensityOfStates::getPseudoUniformKDistribution(int id, int &weight) const
{
	// First we find the unit!
	Unit* unitref = system->getUnit(id);

	// We allocate the k-points
	mat k(3,N_samples);
	
	if(unitref->getBrillouin().n_cols == 2) // A line!
	{
		weight = 2; // We will dived the line into two peices (symmetry).
		
		// The origin!
		vec orig = arma::sum(unitref->getBrillouin(),1)/2;
		// The end!
		vec end = unitref->getBrillouin().col(0); 

		_Cilk_for(long nk = 0; nk < N_samples; nk++)
		{
			double f = (double)nk/(double)N_samples;
			k.col(nk) = (1-f)*orig + f*end;
		}
	}
	else if(unitref->getBrillouin().n_cols == 4) // A square!
	{
		weight = 4; // We will dived the line into four peices (symmetry).
		
		//Note the brillouin points are sorted counter clockwize!

		// The first corner!
		vec p1 = unitref->getBrillouin().col(0);
		vec p2 = (unitref->getBrillouin().col(0)+unitref->getBrillouin().col(1))/2;
		vec p3 = (unitref->getBrillouin().col(0)+unitref->getBrillouin().col(2))/2;
		vec p4 = (unitref->getBrillouin().col(0)+unitref->getBrillouin().col(3))/2;

		// There are two triangles!
		double area1 = arma::norm(arma::cross(p2-p1,p3-p1),2)/2;
		double area2 = arma::norm(arma::cross(p3-p1,p4-p1),2)/2;
		double totalarea = area1 + area2;
		
		// The points are distributed between the two triangles!
		size_t point_count_1 = (int)(area1/totalarea);
		size_t point_count_2 = N_samples-point_count_1;

		// To simulate uniformity, a minimum distance must be met!
		double min_dist_1 = sqrt(area1/point_count_1)/1.5;
		double min_dist_2 = sqrt(area2/point_count_1)/1.5;

		// A random point generator (in triangle)!
		static const auto randomPoint = [] (vec p1, vec p2, vec p3) {
			vec2 f = arma::randu<vec>(2);
			return (1-sqrt(f(1)))*p3 + sqrt(f(1))*((1-f(0))*p1+f(0)*p2);
		};

		// Note that the algorithm can continue infinitely, so we make a max!
		const size_t max_iter = 10000;
		size_t iter = 0;
		
		// The first triangle is made!
		for(long nk = 0; nk < point_count_1; nk++)
		{
			vec p = randomPoint(p1,p2,p3);
			bool allowed = true;
			
			// Note if iter >= max_iter this wil not be checked!
			for(long nnk = 0; nnk <= nk && allowed && iter < max_iter; nnk++)
			{
				if(arma::norm(p-k.col(nnk),2) < min_dist_1)
					allowed = false;
			}

			// Another attempt has been made!
			iter++;

			if(allowed)
				k.col(nk) = p;
			else
				nk--; // Try again!
		}

		// reset iter!
		iter = 0;
		
		// The second triangle is made!
		for(long nk = point_count_1; nk < N_samples; nk++)
		{
			vec p = randomPoint(p1,p3,p4);
			bool allowed = true;
			
			// Note if iter >= max_iter this wil not be checked!
			for(long nnk = point_count_1; nnk <= nk && allowed && iter < max_iter; nnk++)
			{
				if(arma::norm(p-k.col(nnk),2) < min_dist_1)
					allowed = false;
			}

			// Another attempt has been made!
			iter++;

			if(allowed)
				k.col(nk) = p;
			else
				nk--; // Try again!
		}

		// Done!
	}
	else if(unitref->getBrillouin().n_cols == 6) // A hexagon!
	{
		weight = 4; // We will dived the line into four peices (symmetry).
		
		//Note the brillouin points are sorted counter clockwize!

		// The first corner!
		vec p1 = unitref->getBrillouin().col(0);
		vec p2 = unitref->getBrillouin().col(1);
		vec p3 = (unitref->getBrillouin().col(1)+unitref->getBrillouin().col(2))/2;
		vec p4 = (unitref->getBrillouin().col(0)+unitref->getBrillouin().col(3))/2;

		// There are two triangles!
		double area1 = arma::norm(arma::cross(p2-p1,p3-p1),2)/2;
		double area2 = arma::norm(arma::cross(p3-p1,p4-p1),2)/2;
		double totalarea = area1 + area2;
		
		// The points are distributed between the two triangles!
		size_t point_count_1 = (int)(area1/totalarea);
		size_t point_count_2 = N_samples-point_count_1;

		// To simulate uniformity, a minimum distance must be met!
		double min_dist_1 = sqrt(area1/point_count_1)/1.5;
		double min_dist_2 = sqrt(area2/point_count_1)/1.5;

		// Note that the algorithm can continue infinitely, so we make a max!
		const size_t max_iter = 10000;
		size_t iter = 0;
		
		// The first triangle is made!
		for(long nk = 0; nk < point_count_1; nk++)
		{
			vec p = randomPoint(p1,p2,p3);
			bool allowed = true;
			
			// Note if iter >= max_iter this wil not be checked!
			for(long nnk = 0; nnk <= nk && allowed && iter < max_iter; nnk++)
			{
				if(arma::norm(p-k.col(nnk),2) < min_dist_1)
					allowed = false;
			}

			// Another attempt has been made!
			iter++;

			if(allowed)
				k.col(nk) = p;
			else
				nk--; // Try again!
		}

		// reset iter!
		iter = 0;
		
		// The second triangle is made!
		for(long nk = point_count_1; nk < N_samples; nk++)
		{
			vec p = randomPoint(p1,p3,p4);
			bool allowed = true;
			
			// Note if iter >= max_iter this wil not be checked!
			for(long nnk = point_count_1; nnk <= nk && allowed && iter < max_iter; nnk++)
			{
				if(arma::norm(p-k.col(nnk),2) < min_dist_1)
					allowed = false;
			}

			// Another attempt has been made!
			iter++;

			if(allowed)
				k.col(nk) = p;
			else
				nk--; // Try again!
		}

		// Done!
	}

	return k;
}